# coding=utf-8
# Filename:    searcher.py
# Author:      ZENGGUANRONG
# Date:        2023-12-12
# description: 核心检索器

import json, requests, copy
import numpy as np
from loguru import logger
from src.searcher.vec_searcher.vec_searcher import VecSearcher
from src.models.vec_model.vec_model import VectorizeModel


class Searcher:
    def __init__(self, model_path, vec_search_path):
        self.vec_model = VectorizeModel(model_path)
        logger.info("load vec_model done")

        self.vec_searcher = VecSearcher()
        self.vec_searcher.load(vec_search_path)
        logger.info("load vec_searcher done")

    def rank(self, query, recall_result):
        """
        faiss 效果似乎有点问题，所以进行了重排
        :param query:
        :param recall_result:
        :return:
        """
        rank_result = []
        for idx in range(len(recall_result)):
            new_sim = self.vec_model.cos_sim(query, recall_result[idx][1][0])
            rank_item = copy.deepcopy(recall_result[idx])
            rank_item.append(new_sim)
            rank_result.append(copy.deepcopy(rank_item))
        rank_result.sort(key=lambda x: x[3], reverse=True)
        return rank_result

    def search(self, query, nums=3):
        logger.info("request: {}".format(query))

        q_vec = self.vec_model.predict_vec(query).cpu().numpy()

        recall_result = self.vec_searcher.search(q_vec, nums)

        # 做了一个重排
        # rank_result = self.rank(query, recall_result)
        # rank_result = list(filter(lambda x:x[4] > 0.8, rank_result))

        # logger.info("response: {}".format(recall_result))
        return recall_result


if __name__ == "__main__":
    VEC_MODEL_PATH = "/home/guweizheng/models/simcse-chinese-roberta-wwm-ext"
    VEC_INDEX_DATA = "baike_qa_vec_index_2w"
    searcher = Searcher(VEC_MODEL_PATH, VEC_INDEX_DATA)
    q = "什么人不能吃花生"
    res = searcher.search(q)
    for i in res:
        print(i)
        print('----------')
#     [13004, ['花生是生吃好好是熟吃好？ ', {'qid': 'qid_5522945436623969373', 'category': '生活-美食/烹饪', 'title': '花生是生吃好好是熟吃好？ ', 'desc': '', 'answer': '不宜生吃花生 花生含有丰富的植物油、蛋白质和维生素，是人人喜爱的食品。但生吃花生有损于健康。 因为花生含脂肪较多，消化吸收缓慢，大量生吃可以引起消化不良。另外，花生在泥土里生长，常被寄生虫卵污染，生吃容易引起寄生虫病。同时，花生常被鼠类污染，易传播自然疫源性疾病，特别是流行性出血热。因此，花生不宜生吃，最好是煮熟后食用。 \r\n'}], 209.10405]
# ----------
# [3427, ['吃什么样的可以去或', {'qid': 'qid_5732839236678556957', 'category': '生活-生活常识', 'title': '如何用食谱去火吃什么样的蔬菜可以去或 ', 'desc': '吃什么样的可以去或', 'answer': '一、喝莲子汤去心火 \r\u3000\u3000表现症状：分虚实两种，虚火表现为低热、盗汗、心烦、口干等；实火表现为反复、口干、小便短赤、心烦易怒等。 \r食疗法：莲子30克(不去莲心)，桅子15克(用纱布包扎)，加冰糖适量，水煎，吃莲子喝汤。 \r\u3000\u3000二、吃猪肝可去肺火 \r\u3000\u3000表现症状：干咳无痰或痰少而粘、潮热盗汗、手足心热、失眠、舌红。 \r食疗法：猪肝1付，菊花30克(用纱布包好)，共煮至肝熟，吃肝喝汤。 \r\u3000\u3000三、喝绿豆粥去胃火 \r\u3000\u3000表现症状：分虚实两种，虚火表现为轻微咳嗽、饮食量少、便秘、腹胀、舌红、少苔；实火表现为上腹不适、口干口苦、大便干硬。 \r食疗法：石膏粉30克，粳米、绿豆各适量，先用水煎煮石膏，然后过滤去渣，取其清液，再加入粳米、绿豆煮粥食之。 \r\u3000\u3000四、喝梨水去肝火 \r\u3000\u3000表现症状：头痛、头晕、耳鸣、眼干、口苦口臭、两肋胀痛。 \r食疗法：川贝母10克捣碎成末，梨2个，削皮切块，加冰糖适量，清水适量炖服。 \r\u3000\u3000五、吃猪腰去肾火 \r\u3000\u3000表现症状：头晕目眩、耳鸣耳聋、腰脊酸软、潮热盗汗、五心烦躁。 \r食疗法：猪腰2只，枸杞子、山萸肉各15克，共放入砂锅内煮至猪腰子熟，吃猪腰子喝汤。 \r\r雪梨水鱼窝 \r\r\u3000\u3000原料： \r\r\u3000\u3000雪梨4个，水鱼300克，南杏仁50克，葱白、姜、龙眼肉、果皮各少许。 \r\r\u3000\u3000做法： \r\r\u3000\u30001、将水鱼放血，杀好洗干净，用大热水去潺，再斩件，以少许盐、果皮及姜丝略腌。 \r\r\u3000\u30002、雪梨去皮及芯，同南杏、葱、姜、果皮、龙眼肉及清水煲30分钟。 \r\r\u3000\u30003、放入水鱼件煲8分钟至水鱼刚熟，盛碟，汤以少盐调味即成。 \r\r\u3000\u3000贴士： \r\r\u3000\u3000水鱼可找店主代切，买后必须将白色膏状脂肪除去，以去腥味；水鱼必须拖水才可使用；水鱼煲得久都不会变老，但熟透就行了，最好不要煲得太久；雪梨要拣重的，用时要去皮去芯，汤水便不会酸涩。\r\r'}], 251.98729]
# ----------
# [2164, ['应该吃什么', {'qid': 'qid_5774354275070495954', 'category': '健康-五官科-口腔科', 'title': '口腔溃疡应该吃什么 ', 'desc': '应该吃什么', 'answer': '喷口腔粘膜溃疡粉最好了！顶多两天好，或是抹蜂蜜也很管用！'}], 257.4323]
# ----------
